import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestRegressor

#step1.5 预处理

# ----------------------------
# 通用函数与缺失值处理部分
# ----------------------------

# 生成正确的负荷列名（T0000~T2345，每15分钟一个，共96列）
def get_load_columns():
    cols = []
    for hour in range(24):
        for minute in [0, 15, 30, 45]:
            col_name = f'T{hour:02d}{minute:02d}'
            cols.append(col_name)
    return cols[:96]


# 检查缺失值比例是否超过50%
def check_missing_ratio(df, data_name):
    # 计算整体缺失值比例（所有单元格中缺失值的占比）
    total_cells = df.size
    missing_cells = df.isnull().sum().sum()
    missing_ratio = missing_cells / total_cells if total_cells > 0 else 1.0  # 空表视为100%缺失

    print(f"{data_name} 缺失值比例：{missing_ratio:.2%}")
    if missing_ratio > 0.5:
        print(f"错误：{data_name} 缺失值超过50%，无法进行有效填充")
        return False  # 缺失率过高，返回失败
    return True  # 缺失率可接受


# 近10天均值填充函数（无样本时使用）
def fill_with_last_10_days_mean(df, target_col, date_col='YMD'):
    df = df.copy()
    df[date_col] = pd.to_datetime(df[date_col])
    df = df.sort_values(by=date_col).reset_index(drop=True)

    for idx in df[df[target_col].isnull()].index:
        current_date = df.loc[idx, date_col]
        past_10_days = df[(df[date_col] < current_date) & (df[target_col].notnull())].tail(10)

        if len(past_10_days) > 0:
            mean_val = past_10_days[target_col].mean()
        else:
            mean_val = df[target_col].mean() if not pd.isna(df[target_col].mean()) else 0
        df.loc[idx, target_col] = mean_val
    return df


# 随机森林填充函数（无样本时调用近10天均值）
def fill_with_rf(df, target_feature):
    df_copy = df.copy()
    df_copy['Year'] = pd.to_datetime(df_copy['YMD']).dt.year
    df_copy['Month'] = pd.to_datetime(df_copy['YMD']).dt.month
    df_copy['Day'] = pd.to_datetime(df_copy['YMD']).dt.day
    df_copy['Weekday'] = pd.to_datetime(df_copy['YMD']).dt.weekday

    train_data = df_copy[df_copy[target_feature].notnull()]
    if len(train_data) == 0:
        print(f"警告：{target_feature} 无训练样本，使用近10天均值填充")
        return fill_with_last_10_days_mean(df_copy, target_feature)

    features = [col for col in df_copy.columns if col not in [target_feature, 'YMD']]
    X_train, y_train = train_data[features], train_data[target_feature]
    X_test = df_copy[df_copy[target_feature].isnull()][features]

    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    if not X_test.empty:
        df_copy.loc[df_copy[target_feature].isnull(), target_feature] = model.predict(X_test)

    return df_copy.drop(columns=['Year', 'Month', 'Day', 'Weekday'])


# 气象数据缺失值处理（保留两位小数）
def handle_meteorological_missing(weather_data):
    # 填充逻辑不变，已在主函数中通过缺失率检查
    print("1. 填充最高温度...")
    weather_data = fill_with_rf(weather_data, 'Max_Temperature').round(2)

    print("2. 填充最低温度...")
    weather_data = fill_with_rf(weather_data, 'Min_Temperature').round(2)

    print("3. 计算平均温度...")
    missing_avg = weather_data['Avg_Temperature'].isnull()
    weather_data.loc[missing_avg, 'Avg_Temperature'] = (
                                                               weather_data.loc[missing_avg, 'Max_Temperature'] +
                                                               weather_data.loc[missing_avg, 'Min_Temperature']
                                                       ) / 2
    weather_data['Avg_Temperature'] = weather_data['Avg_Temperature'].round(2)

    print("4. 填充平均湿度...")
    weather_data = fill_with_rf(weather_data, 'Avg_Humidity').round(2)

    print("5. 填充降雨量...")
    weather_data = fill_with_rf(weather_data, 'Rainfall').round(2)

    return weather_data


# 负荷数据缺失值处理（保留两位小数）
def handle_load_missing(load_data):
    # 填充逻辑不变，已在主函数中通过缺失率检查
    load_cols = get_load_columns()
    for col in load_cols:
        load_data = fill_with_last_10_days_mean(load_data, col)
        load_data[col] = load_data[col].round(2)
    return load_data


# ----------------------------
# 异常值处理逻辑
# ----------------------------

# 负荷数据异常值检测函数
def check_negative_or_zero(df, cols):
    return df[cols].lt(0).any(axis=1) | df[cols].eq(0).any(axis=1)


def check_peak_in_valley(df, cols):
    daily_mean = df[cols].mean(axis=1)
    valley_cols = [col for col in cols if 'T0200' <= col <= 'T0545']
    return df[valley_cols].gt(daily_mean * 3).any(axis=1)


def check_sudden_jump(df, cols):
    load_matrix = df[cols].values
    load_matrix = np.where(load_matrix == 0, 1e-6, load_matrix)
    diff_ratio = np.abs(np.diff(load_matrix, axis=1)) / load_matrix[:, :-1]
    return (diff_ratio > 0.5).any(axis=1)


def detect_weekday_deviation(df, cols):
    df['weekday'] = pd.to_datetime(df['YMD']).dt.weekday
    weekday_means = df.groupby('weekday')[cols].mean()
    deviations = df[cols].sub(df['weekday'].map(weekday_means.mean(axis=1)), axis=0).abs().div(
        df['weekday'].map(weekday_means.std(axis=1)) + 1e-6,
        axis=0
    )
    return (deviations > 3).any(axis=1)


def detect_outliers_iqr(df, cols):
    q1 = df[cols].quantile(0.25)
    q3 = df[cols].quantile(0.75)
    iqr = q3 - q1
    lower_bound = q1 - 1.5 * iqr
    upper_bound = q3 + 1.5 * iqr
    return ((df[cols] < lower_bound) | (df[cols] > upper_bound)).any(axis=1)


# 负荷数据异常值处理（替换异常值+删除严重异常）
def replace_outliers_with_q95(df, cols):
    for col in cols:
        q95 = df[col].quantile(0.95)
        q1 = df[col].quantile(0.25)
        q3 = df[col].quantile(0.75)
        upper_bound = q3 + 1.5 * (q3 - q1)
        severe_condition = df[col] > upper_bound * 1.2
        mild_condition = df[col] > upper_bound
        df[col] = np.where(mild_condition & ~severe_condition, q95, df[col])
    return df


# 负荷数据异常值处理主函数
def clean_load_data(load_data):
    load_cols = get_load_columns()

    # 添加每日统计指标
    load_data['MaxLoad'] = load_data[load_cols].max(axis=1)
    load_data['MinLoad'] = load_data[load_cols].min(axis=1)
    load_data['MeanLoad'] = load_data[load_cols].mean(axis=1)
    print("\n添加统计指标后的数据：")
    print(load_data.head())

    # 可视化原始负荷趋势
    plt.figure(figsize=(12, 6))
    plt.plot(load_data['MaxLoad'], marker='*', linestyle='None', label='MaxLoad')
    plt.plot(load_data['MinLoad'], marker='*', linestyle='None', label='MinLoad')
    plt.plot(load_data['MeanLoad'], marker='*', linestyle='None', label='MeanLoad')
    plt.legend()
    plt.title("Daily Load Trends (Before Outlier Handling)")
    plt.show()

    # 箱线图查看原始分布
    plt.figure(figsize=(18, 6))
    sns.boxplot(data=load_data[load_cols])
    plt.xticks(rotation=45)
    plt.title("Boxplot of Raw Load Data")
    plt.show()

    # 多维度异常检测
    negative_mask = check_negative_or_zero(load_data, load_cols)
    peak_in_valley_mask = check_peak_in_valley(load_data, load_cols)
    sudden_jump_mask = check_sudden_jump(load_data, load_cols)
    weekday_deviation_mask = detect_weekday_deviation(load_data, load_cols)
    stat_mask = detect_outliers_iqr(load_data, load_cols)

    # 综合异常标记
    total_outlier_mask = negative_mask | peak_in_valley_mask | sudden_jump_mask | weekday_deviation_mask | stat_mask
    print(f"\n检测到异常数据行数：{total_outlier_mask.sum()}")

    # 替换轻度异常值
    load_data = replace_outliers_with_q95(load_data, load_cols)

    # 删除严重异常记录
    cleaned_data = load_data[~total_outlier_mask].copy()

    # 更新统计指标
    cleaned_data['MaxLoad'] = cleaned_data[load_cols].max(axis=1)
    cleaned_data['MinLoad'] = cleaned_data[load_cols].min(axis=1)
    cleaned_data['MeanLoad'] = cleaned_data[load_cols].mean(axis=1)

    # 可视化处理后趋势
    plt.figure(figsize=(12, 6))
    plt.plot(cleaned_data['MaxLoad'], marker='*', linestyle='None', label='MaxLoad')
    plt.plot(cleaned_data['MinLoad'], marker='*', linestyle='None', label='MinLoad')
    plt.plot(cleaned_data['MeanLoad'], marker='*', linestyle='None', label='MeanLoad')
    plt.legend()
    plt.title("Load Trends After Outlier Handling")
    plt.show()

    return cleaned_data.round(2)


# 气象数据异常值检测函数
def check_temperature_range(df):
    temp_cols = ['Max_Temperature', 'Min_Temperature', 'Avg_Temperature']
    return df[temp_cols].lt(-30).any(axis=1) | df[temp_cols].gt(45).any(axis=1)


def check_temp_order(df):
    return (df['Max_Temperature'] < df['Avg_Temperature']) | (df['Avg_Temperature'] < df['Min_Temperature'])


def check_temp_diff(df):
    return (df['Max_Temperature'] - df['Min_Temperature']) > 25


def check_humidity_range(df):
    return (df['Avg_Humidity'] < 0) | (df['Avg_Humidity'] > 100)


def check_rainfall(df):
    return (df['Rainfall'] < 0) | (df['Rainfall'] > 200)


# 气象数据异常值处理（删除异常记录）
def clean_meteorological_data(weather_data):
    temp_mask = check_temperature_range(weather_data)
    order_mask = check_temp_order(weather_data)
    diff_mask = check_temp_diff(weather_data)
    humidity_mask = check_humidity_range(weather_data)
    rainfall_mask = check_rainfall(weather_data)

    weather_outlier_mask = temp_mask | order_mask | diff_mask | humidity_mask | rainfall_mask
    print(f"检测到气象异常数据行数：{weather_outlier_mask.sum()}")

    return weather_data[~weather_outlier_mask].copy().round(2)


# ----------------------------
# 主函数（添加缺失率检查）
# ----------------------------

def main():
    try:
        print("加载原始数据...")
        load_data = pd.read_csv('data/Load.csv')
        weather_data = pd.read_csv('data/Weather.csv')

        # 第一步：检查缺失值比例（超过50%则终止）
        print("\n===== 检查缺失值比例 =====")
        load_valid = check_missing_ratio(load_data, "负荷数据")
        weather_valid = check_missing_ratio(weather_data, "气象数据")

        if not (load_valid and weather_valid):
            print("预测失败：数据缺失过多，无法继续处理")
            return  # 缺失率过高，终止程序

        # 第二步：缺失值处理（通过检查后才执行）
        print("\n===== 缺失值处理（保留两位小数） =====")
        filled_weather = handle_meteorological_missing(weather_data)
        filled_load = handle_load_missing(load_data)

        filled_weather.to_csv('data/Area1_Weather_missing_filled.csv', index=False)
        filled_load.to_csv('data/Area1_Load_missing_filled.csv', index=False)
        print("\n缺失值处理完成，中间文件已保存")

        # 第三步：异常值处理
        print("\n===== 异常值处理 =====")
        cleaned_weather = clean_meteorological_data(filled_weather)
        cleaned_load = clean_load_data(filled_load)

        cleaned_weather.to_csv('data/Area1_Weather_clean.csv', index=False)
        cleaned_load.to_csv('data/Area1_Load_clean.csv', index=False)
        print("\n所有处理完成，最终文件已保存（两位小数）")

    except Exception as e:
        print(f"处理异常：{str(e)}")


if __name__ == '__main__':
    main()